Skip to content

Commit 405f28e

Browse files
VinciGit00VinciGit00
authored
Merge pull request #606 from Santabot123/pre/beta
Screenshot scraping
2 parents d7f6036 + 735120d commit 405f28e

File tree

3 files changed

+249
-0
lines changed

3 files changed

+249
-0
lines changed

‎scrapegraphai/screenshot_scraping/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
"""
2+
__init__.py file for screenshot_scraping folder
3+
"""
4+
5+
6+
from .screenshot_preparation import take_screenshot, select_area_with_opencv, select_area_with_ipywidget, crop_image
7+
from .text_detection import detect_text
8+

‎scrapegraphai/screenshot_scraping/screenshot_preparation.py

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
import asyncio
2+
from playwright.async_api import async_playwright
3+
4+
from io import BytesIO
5+
from PIL import Image, ImageGrab
6+
7+
8+
async def take_screenshot(url: str, save_path: str = None, quality: int = 100):
9+
"""
10+
Takes a screenshot of a webpage at the specified URL and saves it if the save_path is specified.
11+
Parameters:
12+
url (str): The URL of the webpage to take a screenshot of.
13+
save_path (str): The path to save the screenshot to. Defaults to None.
14+
quality (int): The quality of the jpeg image, between 1 and 100. Defaults to 100.
15+
Returns:
16+
PIL.Image: The screenshot of the webpage as a PIL Image object.
17+
"""
18+
19+
async with async_playwright() as p:
20+
browser = await p.chromium.launch(headless=True)
21+
page = await browser.new_page()
22+
await page.goto(url)
23+
image_bytes = await page.screenshot(path=save_path, type="jpeg", full_page=True, quality=quality)
24+
await browser.close()
25+
return Image.open(BytesIO(image_bytes))
26+
27+
28+
def select_area_with_opencv(image):
29+
"""
30+
Allows you to manually select an image area using OpenCV. It is recommended to use this function if your project is on your computer, otherwise use select_area_with_ipywidget().
31+
Parameters:
32+
image (PIL.Image): The image from which to select an area.
33+
Returns:
34+
A tuple containing the LEFT, TOP, RIGHT, and BOTTOM coordinates of the selected area.
35+
"""
36+
37+
import cv2 as cv
38+
import numpy as np
39+
40+
fullscreen_screenshot = ImageGrab.grab()
41+
dw, dh = fullscreen_screenshot.size
42+
43+
def draw_selection_rectanlge(event, x, y, flags, param):
44+
global ix, iy, drawing, overlay, img
45+
if event == cv.EVENT_LBUTTONDOWN:
46+
drawing = True
47+
ix, iy = x, y
48+
elif event == cv.EVENT_MOUSEMOVE:
49+
if drawing == True:
50+
cv.rectangle(img, (ix, iy), (x, y), (41, 215, 162), -1)
51+
cv.putText(img, 'PRESS ANY KEY TO SELECT THIS AREA', (ix,
52+
iy-10), cv.FONT_HERSHEY_SIMPLEX, 1.5, (55, 46, 252), 5)
53+
img = cv.addWeighted(overlay, alpha, img, 1 - alpha, 0)
54+
elif event == cv.EVENT_LBUTTONUP:
55+
global LEFT, TOP, RIGHT, BOTTOM
56+
57+
drawing = False
58+
if ix < x:
59+
LEFT = int(ix)
60+
RIGHT = int(x)
61+
else:
62+
LEFT = int(x)
63+
RIGHT = int(ix)
64+
if iy < y:
65+
TOP = int(iy)
66+
BOTTOM = int(y)
67+
else:
68+
TOP = int(y)
69+
BOTTOM = int(iy)
70+
71+
global drawing, ix, iy, overlay, img
72+
drawing = False
73+
ix, iy = -1, -1
74+
75+
img = np.array(image)
76+
img = cv.cvtColor(img, cv.COLOR_RGB2BGR)
77+
78+
img = cv.rectangle(
79+
img, (0, 0), (image.size[0], image.size[1]), (0, 0, 255), 10)
80+
img = cv.putText(img, 'SELECT AN AREA', (int(
81+
image.size[0]*0.3), 100), cv.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 5)
82+
83+
overlay = img.copy()
84+
alpha = 0.3
85+
86+
while True:
87+
cv.namedWindow('SELECT AREA', cv.WINDOW_KEEPRATIO)
88+
cv.setMouseCallback('SELECT AREA', draw_selection_rectanlge)
89+
cv.resizeWindow('SELECT AREA', int(
90+
image.size[0]/(image.size[1]/dh)), dh)
91+
92+
cv.imshow('SELECT AREA', img)
93+
94+
if cv.waitKey(20) > -1:
95+
break
96+
97+
cv.destroyAllWindows()
98+
return LEFT, TOP, RIGHT, BOTTOM
99+
100+
101+
def select_area_with_ipywidget(image):
102+
"""
103+
Allows you to manually select an image area using ipywidgets. It is recommended to use this function if your project is in Google Colab, Kaggle or other similar platform, otherwise use select_area_with_opencv().
104+
Parameters:
105+
image (PIL Image): The input image.
106+
Returns:
107+
None
108+
"""
109+
110+
import matplotlib.pyplot as plt
111+
import numpy as np
112+
from ipywidgets import interact, IntSlider
113+
import ipywidgets as widgets
114+
from PIL import Image
115+
116+
img_array = np.array(image)
117+
118+
print(img_array.shape)
119+
120+
def update_plot(top_bottom, left_right, image_size):
121+
plt.figure(figsize=(image_size, image_size))
122+
plt.imshow(img_array)
123+
plt.axvline(x=left_right[0], color='blue', linewidth=1)
124+
plt.text(left_right[0]+1, -25, 'LEFT', rotation=90, color='blue')
125+
plt.axvline(x=left_right[1], color='red', linewidth=1)
126+
plt.text(left_right[1]+1, -25, 'RIGHT', rotation=90, color='red')
127+
128+
plt.axhline(y=img_array.shape[0] -
129+
top_bottom[0], color='green', linewidth=1)
130+
plt.text(-100, img_array.shape[0] -
131+
top_bottom[0]+1, 'BOTTOM', color='green')
132+
plt.axhline(y=img_array.shape[0]-top_bottom[1],
133+
color='darkorange', linewidth=1)
134+
plt.text(-100, img_array.shape[0] -
135+
top_bottom[1]+1, 'TOP', color='darkorange')
136+
plt.axis('off')
137+
plt.show()
138+
139+
top_bottom_slider = widgets.IntRangeSlider(
140+
value=[int(img_array.shape[0]*0.25), int(img_array.shape[0]*0.75)],
141+
min=0,
142+
max=img_array.shape[0],
143+
step=1,
144+
description='top_bottom:',
145+
disabled=False,
146+
continuous_update=True,
147+
orientation='vertical',
148+
readout=True,
149+
readout_format='d',
150+
)
151+
152+
left_right_slider = widgets.IntRangeSlider(
153+
value=[int(img_array.shape[1]*0.25), int(img_array.shape[1]*0.75)],
154+
min=0,
155+
max=img_array.shape[1],
156+
step=1,
157+
description='left_right:',
158+
disabled=False,
159+
continuous_update=True,
160+
orientation='horizontal',
161+
readout=True,
162+
readout_format='d',
163+
)
164+
image_size_bt = widgets.BoundedIntText(
165+
value=10,
166+
min=2,
167+
max=20,
168+
step=1,
169+
description='Image size:',
170+
disabled=False
171+
)
172+
173+
interact(update_plot, top_bottom=top_bottom_slider,
174+
left_right=left_right_slider, image_size=image_size_bt)
175+
176+
return left_right_slider, top_bottom_slider
177+
178+
179+
def crop_image(image, LEFT=None, TOP=None, RIGHT=None, BOTTOM=None, save_path: str = None):
180+
"""
181+
Crop an image using the specified coordinates.
182+
Parameters:
183+
image (PIL.Image): The image to be cropped.
184+
LEFT (int, optional): The x-coordinate of the left edge of the crop area. Defaults to None.
185+
TOP (int, optional): The y-coordinate of the top edge of the crop area. Defaults to None.
186+
RIGHT (int, optional): The x-coordinate of the right edge of the crop area. Defaults to None.
187+
BOTTOM (int, optional): The y-coordinate of the bottom edge of the crop area. Defaults to None.
188+
save_path (str, optional): The path to save the cropped image. Defaults to None.
189+
Returns:
190+
PIL.Image: The cropped image.
191+
Notes:
192+
If any of the coordinates (LEFT, TOP, RIGHT, BOTTOM) is None, it will be set to the corresponding edge of the image.
193+
If save_path is specified, the cropped image will be saved as a JPEG file at the specified path.
194+
"""
195+
196+
if LEFT is None:
197+
LEFT = 0
198+
if TOP is None:
199+
TOP = 0
200+
if RIGHT is None:
201+
RIGHT = image.size[0]
202+
if BOTTOM is None:
203+
BOTTOM = image.size[1]
204+
205+
croped_image = image.crop((LEFT, TOP, RIGHT, BOTTOM))
206+
if save_path is not None:
207+
from pathlib import Path
208+
croped_image.save(save_path, "JPEG")
209+
210+
return image.crop((LEFT, TOP, RIGHT, BOTTOM))
211+
212+

‎scrapegraphai/screenshot_scraping/text_detection.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from surya.ocr import run_ocr
2+
import numpy as np
3+
from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
4+
from surya.model.recognition.model import load_model as load_rec_model
5+
from surya.model.recognition.processor import load_processor as load_rec_processor
6+
7+
8+
def detect_text(image, languages: list = ["en"]):
9+
"""
10+
Detects and extracts text from a given image.
11+
Parameters:
12+
image (PIL Image): The input image to extract text from.
13+
lahguages (list): A list of languages to detect text in. Defaults to ["en"]. List of languages can be found here: https://github.com/VikParuchuri/surya/blob/master/surya/languages.py
14+
Returns:
15+
str: The extracted text from the image.
16+
Notes:
17+
Model weights will automatically download the first time you run this function.
18+
"""
19+
20+
langs = languages
21+
det_processor, det_model = load_det_processor(), load_det_model()
22+
rec_model, rec_processor = load_rec_model(), load_rec_processor()
23+
predictions = run_ocr([image], [langs], det_model,
24+
det_processor, rec_model, rec_processor)
25+
26+
text = "\n".join([line.text for line in predictions[0].text_lines])
27+
return text
28+
29+

0 commit comments

Comments
 (0)