mirror of
https://github.com/VinciGit00/Scrapegraph-ai.git
synced 2026-07-01 21:00:48 +08:00
add __init__.py and docstrings
This commit is contained in:
parent
e11f0cd28d
commit
b4f8ea420b
8
scrapegraphai/sreenshot_scraping/__init__.py
Normal file
8
scrapegraphai/sreenshot_scraping/__init__.py
Normal file
@ -0,0 +1,8 @@
|
||||
"""
|
||||
__init__.py file for screenshot_scraping folder
|
||||
"""
|
||||
|
||||
|
||||
from .screenshot_preparation import take_screenshot, select_area_with_opencv, select_area_with_ipywidget, crop_image
|
||||
from .text_detection import detect_text
|
||||
|
||||
@ -3,116 +3,140 @@ from playwright.async_api import async_playwright
|
||||
|
||||
from io import BytesIO
|
||||
from PIL import Image, ImageGrab
|
||||
import os
|
||||
|
||||
|
||||
async def take_screenshot(url:str, image_name:str=None, quality:int=100):
|
||||
async def take_screenshot(url: str, save_path: str = None, quality: int = 100):
|
||||
"""
|
||||
Takes a screenshot of a webpage at the specified URL and saves it if the save_path is specified.
|
||||
Parameters:
|
||||
url (str): The URL of the webpage to take a screenshot of.
|
||||
save_path (str): The path to save the screenshot to. Defaults to None.
|
||||
quality (int): The quality of the jpeg image, between 1 and 100. Defaults to 100.
|
||||
Returns:
|
||||
PIL.Image: The screenshot of the webpage as a PIL Image object.
|
||||
"""
|
||||
|
||||
async with async_playwright() as p:
|
||||
browser = await p.chromium.launch(headless=True)
|
||||
page = await browser.new_page()
|
||||
await page.goto(url)
|
||||
save_path=None
|
||||
if image_name is not None:
|
||||
dir_path = os.path.abspath('')
|
||||
save_path = dir_path + "/scrapegraphai/sreenshot_scraping/saved_screenshots/" + image_name
|
||||
image_bytes = await page.screenshot(path=save_path,type="jpeg",full_page=True,quality=quality)
|
||||
image_bytes = await page.screenshot(path=save_path, type="jpeg", full_page=True, quality=quality)
|
||||
await browser.close()
|
||||
return Image.open(BytesIO(image_bytes))
|
||||
|
||||
|
||||
def select_area_with_opencv(image):
|
||||
"""
|
||||
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().
|
||||
Parameters:
|
||||
image (PIL.Image): The image from which to select an area.
|
||||
Returns:
|
||||
A tuple containing the LEFT, TOP, RIGHT, and BOTTOM coordinates of the selected area.
|
||||
"""
|
||||
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
|
||||
fullscreen_screenshot = ImageGrab.grab()
|
||||
dw, dh=fullscreen_screenshot.size
|
||||
dw, dh = fullscreen_screenshot.size
|
||||
|
||||
def draw_selection_rectanlge(event, x, y, flags, param):
|
||||
global ix,iy,drawing,overlay ,img
|
||||
global ix, iy, drawing, overlay, img
|
||||
if event == cv.EVENT_LBUTTONDOWN:
|
||||
drawing = True
|
||||
ix,iy = x,y
|
||||
ix, iy = x, y
|
||||
elif event == cv.EVENT_MOUSEMOVE:
|
||||
if drawing == True:
|
||||
cv.rectangle(img, (ix, iy), (x, y), (41, 215, 162), -1)
|
||||
cv.putText(img, 'PRESS ANY KEY TO SELECT THIS AREA', (ix, iy-10), cv.FONT_HERSHEY_SIMPLEX, 1.5, (55,46,252), 5)
|
||||
img=cv.addWeighted(overlay, alpha, img, 1 - alpha, 0)
|
||||
cv.putText(img, 'PRESS ANY KEY TO SELECT THIS AREA', (ix,
|
||||
iy-10), cv.FONT_HERSHEY_SIMPLEX, 1.5, (55, 46, 252), 5)
|
||||
img = cv.addWeighted(overlay, alpha, img, 1 - alpha, 0)
|
||||
elif event == cv.EVENT_LBUTTONUP:
|
||||
global LEFT,TOP,RIGHT,BOTTOM
|
||||
|
||||
global LEFT, TOP, RIGHT, BOTTOM
|
||||
|
||||
drawing = False
|
||||
if ix<x:
|
||||
LEFT=int(ix)
|
||||
RIGHT=int(x)
|
||||
if ix < x:
|
||||
LEFT = int(ix)
|
||||
RIGHT = int(x)
|
||||
else:
|
||||
LEFT=int(x)
|
||||
RIGHT=int(ix)
|
||||
if iy<y:
|
||||
TOP=int(iy)
|
||||
BOTTOM=int(y)
|
||||
LEFT = int(x)
|
||||
RIGHT = int(ix)
|
||||
if iy < y:
|
||||
TOP = int(iy)
|
||||
BOTTOM = int(y)
|
||||
else:
|
||||
TOP=int(y)
|
||||
BOTTOM=int(iy)
|
||||
TOP = int(y)
|
||||
BOTTOM = int(iy)
|
||||
|
||||
global drawing,ix,iy,overlay,img
|
||||
global drawing, ix, iy, overlay, img
|
||||
drawing = False
|
||||
ix,iy = -1,-1
|
||||
ix, iy = -1, -1
|
||||
|
||||
img =np.array(image)
|
||||
img = np.array(image)
|
||||
img = cv.cvtColor(img, cv.COLOR_RGB2BGR)
|
||||
|
||||
img=cv.rectangle(img, (0, 0), (image.size[0], image.size[1]), (0,0,255), 10)
|
||||
img=cv.putText(img, 'SELECT AN AREA', (int(image.size[0]*0.3), 100), cv.FONT_HERSHEY_SIMPLEX, 2, (0,0,255), 5)
|
||||
|
||||
img = cv.rectangle(
|
||||
img, (0, 0), (image.size[0], image.size[1]), (0, 0, 255), 10)
|
||||
img = cv.putText(img, 'SELECT AN AREA', (int(
|
||||
image.size[0]*0.3), 100), cv.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 5)
|
||||
|
||||
overlay = img.copy()
|
||||
alpha = 0.3
|
||||
|
||||
while True:
|
||||
cv.namedWindow('SELECT AREA', cv.WINDOW_KEEPRATIO)
|
||||
cv.namedWindow('SELECT AREA', cv.WINDOW_KEEPRATIO)
|
||||
cv.setMouseCallback('SELECT AREA', draw_selection_rectanlge)
|
||||
cv.resizeWindow('SELECT AREA', int(image.size[0]/(image.size[1]/dh)), dh)
|
||||
|
||||
cv.imshow('SELECT AREA',img)
|
||||
cv.resizeWindow('SELECT AREA', int(
|
||||
image.size[0]/(image.size[1]/dh)), dh)
|
||||
|
||||
cv.imshow('SELECT AREA', img)
|
||||
|
||||
if cv.waitKey(20) > -1:
|
||||
break
|
||||
|
||||
|
||||
cv.destroyAllWindows()
|
||||
return LEFT, TOP, RIGHT, BOTTOM
|
||||
|
||||
|
||||
def select_area_with_ipywidget(image):
|
||||
"""
|
||||
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().
|
||||
Parameters:
|
||||
image (PIL Image): The input image.
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from ipywidgets import interact, IntSlider
|
||||
import ipywidgets as widgets
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# Load an image
|
||||
# image_path = 'piping-yes-when-running-scripts-from-curl.jpeg' # Replace with your image path
|
||||
# image = Image.open(image_path)
|
||||
img_array = np.array(image)
|
||||
|
||||
print(img_array.shape)
|
||||
|
||||
def update_plot(top_bottom,left_right,image_size):
|
||||
plt.figure(figsize = (image_size,image_size))
|
||||
def update_plot(top_bottom, left_right, image_size):
|
||||
plt.figure(figsize=(image_size, image_size))
|
||||
plt.imshow(img_array)
|
||||
plt.axvline(x=left_right[0], color='blue', linewidth=1)
|
||||
plt.text(left_right[0]+1,-25,'LEFT',rotation=90,color='blue')
|
||||
plt.text(left_right[0]+1, -25, 'LEFT', rotation=90, color='blue')
|
||||
plt.axvline(x=left_right[1], color='red', linewidth=1)
|
||||
plt.text(left_right[1]+1,-25,'RIGHT',rotation=90, color='red')
|
||||
plt.text(left_right[1]+1, -25, 'RIGHT', rotation=90, color='red')
|
||||
|
||||
|
||||
plt.axhline(y=img_array.shape[0]-top_bottom[0], color='green', linewidth=1)
|
||||
plt.text(-100,img_array.shape[0]-top_bottom[0]+1,'BOTTOM', color='green')
|
||||
plt.axhline(y=img_array.shape[0]-top_bottom[1], color='darkorange', linewidth=1)
|
||||
plt.text(-100,img_array.shape[0]-top_bottom[1]+1,'TOP', color='darkorange')
|
||||
plt.axis('off') # Hide axes
|
||||
plt.axhline(y=img_array.shape[0] -
|
||||
top_bottom[0], color='green', linewidth=1)
|
||||
plt.text(-100, img_array.shape[0] -
|
||||
top_bottom[0]+1, 'BOTTOM', color='green')
|
||||
plt.axhline(y=img_array.shape[0]-top_bottom[1],
|
||||
color='darkorange', linewidth=1)
|
||||
plt.text(-100, img_array.shape[0] -
|
||||
top_bottom[1]+1, 'TOP', color='darkorange')
|
||||
plt.axis('off')
|
||||
plt.show()
|
||||
|
||||
|
||||
top_bottom_slider=widgets.IntRangeSlider(
|
||||
top_bottom_slider = widgets.IntRangeSlider(
|
||||
value=[int(img_array.shape[0]*0.25), int(img_array.shape[0]*0.75)],
|
||||
min=0,
|
||||
max=img_array.shape[0],
|
||||
@ -125,7 +149,7 @@ def select_area_with_ipywidget(image):
|
||||
readout_format='d',
|
||||
)
|
||||
|
||||
left_right_slider=widgets.IntRangeSlider(
|
||||
left_right_slider = widgets.IntRangeSlider(
|
||||
value=[int(img_array.shape[1]*0.25), int(img_array.shape[1]*0.75)],
|
||||
min=0,
|
||||
max=img_array.shape[1],
|
||||
@ -137,7 +161,7 @@ def select_area_with_ipywidget(image):
|
||||
readout=True,
|
||||
readout_format='d',
|
||||
)
|
||||
image_size_bt=widgets.BoundedIntText(
|
||||
image_size_bt = widgets.BoundedIntText(
|
||||
value=10,
|
||||
min=2,
|
||||
max=20,
|
||||
@ -146,11 +170,44 @@ def select_area_with_ipywidget(image):
|
||||
disabled=False
|
||||
)
|
||||
|
||||
interact(update_plot,top_bottom= top_bottom_slider ,left_right=left_right_slider ,image_size=image_size_bt)
|
||||
interact(update_plot, top_bottom=top_bottom_slider,
|
||||
left_right=left_right_slider, image_size=image_size_bt)
|
||||
|
||||
|
||||
def crop_image(image, LEFT=None, TOP=None, RIGHT=None, BOTTOM=None, save_path: str = None):
|
||||
"""
|
||||
Crop an image using the specified coordinates.
|
||||
Parameters:
|
||||
image (PIL.Image): The image to be cropped.
|
||||
LEFT (int, optional): The x-coordinate of the left edge of the crop area. Defaults to None.
|
||||
TOP (int, optional): The y-coordinate of the top edge of the crop area. Defaults to None.
|
||||
RIGHT (int, optional): The x-coordinate of the right edge of the crop area. Defaults to None.
|
||||
BOTTOM (int, optional): The y-coordinate of the bottom edge of the crop area. Defaults to None.
|
||||
save_path (str, optional): The path to save the cropped image. Defaults to None.
|
||||
Returns:
|
||||
PIL.Image: The cropped image.
|
||||
Notes:
|
||||
If any of the coordinates (LEFT, TOP, RIGHT, BOTTOM) is None, it will be set to the corresponding edge of the image.
|
||||
If save_path is specified, the cropped image will be saved as a JPEG file at the specified path.
|
||||
"""
|
||||
|
||||
if LEFT is None:
|
||||
LEFT = 0
|
||||
if TOP is None:
|
||||
TOP = 0
|
||||
if RIGHT is None:
|
||||
RIGHT = image.size[0]
|
||||
if BOTTOM is None:
|
||||
BOTTOM = image.size[1]
|
||||
|
||||
croped_image = image.crop((LEFT, TOP, RIGHT, BOTTOM))
|
||||
if save_path is not None:
|
||||
from pathlib import Path
|
||||
croped_image.save(save_path, "JPEG")
|
||||
|
||||
def crop_image(image,TOP=None, BOTTOM=None, LEFT=None, RIGHT=None, save:bool=True):
|
||||
return image.crop((LEFT, TOP, RIGHT, BOTTOM))
|
||||
|
||||
|
||||
# image=asyncio.run(take_screenshot("https://unix.stackexchange.com/questions/690233/piping-yes-when-running-scripts-from-curl"))
|
||||
|
||||
# image=asyncio.run(take_screenshot("https://unix.stackexchange.com/questions/690233/piping-yes-when-running-scripts-from-curl", save_path="saved_screenshots/image.jpeg"))
|
||||
# print(select_area_with_opencv(image))
|
||||
@ -1,4 +1,3 @@
|
||||
from PIL import Image
|
||||
from surya.ocr import run_ocr
|
||||
import numpy as np
|
||||
from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
|
||||
@ -6,14 +5,24 @@ from surya.model.recognition.model import load_model as load_rec_model
|
||||
from surya.model.recognition.processor import load_processor as load_rec_processor
|
||||
|
||||
|
||||
def DetectText(image, lahguages:list = ["en"]):
|
||||
def detect_text(image, lahguages: list = ["en"]):
|
||||
"""
|
||||
Detects and extracts text from a given image.
|
||||
Parameters:
|
||||
image (PIL Image): The input image to extract text from.
|
||||
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
|
||||
Returns:
|
||||
str: The extracted text from the image.
|
||||
Notes:
|
||||
Model weights will automatically download the first time you run this function.
|
||||
"""
|
||||
|
||||
# image = Image.fromarray(image_array.astype('uint8'), 'RGB')
|
||||
langs = lahguages
|
||||
det_processor, det_model = load_det_processor(), load_det_model()
|
||||
rec_model, rec_processor = load_rec_model(), load_rec_processor()
|
||||
predictions = run_ocr([image], [langs], det_model, det_processor, rec_model, rec_processor)
|
||||
predictions = run_ocr([image], [langs], det_model,
|
||||
det_processor, rec_model, rec_processor)
|
||||
text = "\n".join([line.text for line in predictions[0].text_lines])
|
||||
return text
|
||||
|
||||
# print(DetectText(Image.open("scrapegraphai\sreenshot_scraping\saved_screenshots\image.jpeg").crop((936,1967,1243,2246))))
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user