Skip to content

Commit

Permalink
up
Browse files Browse the repository at this point in the history
  • Loading branch information
guoerjun committed Aug 16, 2024
1 parent 96b3cfb commit 2b270da
Show file tree
Hide file tree
Showing 9 changed files with 159 additions and 24 deletions.
8 changes: 7 additions & 1 deletion detectron/demo/event_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from sam_everything import SamAnything
from retriever import knowledgeBase
from yolo_predictor import YOLOPredictor
from ocr import do_ocr
params = {
"algo_type": None,
"input_image":None
Expand Down Expand Up @@ -110,6 +111,10 @@ def create_event_handlers():
do_sam_everything,gradio('sam_input','sam_video_input','sam_version'),gradio("sam_output","sam_video_output")
)

components["submit_ocr_btn"].click(
do_ocr,gradio('ocr_type','ocr_input'),gradio("ocr_output","ocr_json_output")
)

components["db_submit_btn"].click(
file_handler,gradio('file_upload','db_name'),gradio("db_view",'db_select','db_test_select')
)
Expand Down Expand Up @@ -254,12 +259,13 @@ def do_sam_version_chage(value):


def do_sam_everything(im,video,version):
print(im)
sam_anything = None
if version == "1":
sam_anything = SamAnything()
elif version == "2":
sam_anything = None
print(im)

image_pil = im['image']
points = im['points']
images = None
Expand Down
Binary file modified detectron/demo/examples/gym.mp4
Binary file not shown.
Binary file added detectron/demo/examples/json.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added detectron/demo/examples/text.jpeg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified detectron/demo/examples/traffic.mp4
Binary file not shown.
61 changes: 61 additions & 0 deletions detectron/demo/ocr.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@

from detectron2.data.detection_utils import read_image,pil_image_to_numpy
from detectron2.utils.visualizer import Visualizer
from sam_everything import visimage_to_pil
import numpy as np
def do_ocr(ocr_type,input):
print(ocr_type)
result = None
np_image = pil_image_to_numpy(input)
if ocr_type == "OCR":
from paddleocr import PaddleOCR
ocr = PaddleOCR(lang='ch', use_angle_cls=True)
# img_path = 'exp.jpeg'
result = ocr.ocr(np_image)
print(result)
result = parse_paddle_result(result)

elif ocr_type == "Easy":
import easyocr
reader = easyocr.Reader(['en','ch_sim']) # 初始化 EasyOCR,选择需要支持的语言(例如英文)
result = reader.readtext(np_image)
result = parse_esay_result(result)

view = Visualizer(np_image)
for item in result:
polygon = np.array(item['box'])
view.draw_polygon(polygon, "k")

vis_image = view.get_output()
pil_images = visimage_to_pil([vis_image])
return pil_images[0],result

def parse_esay_result(data):
results = []
for entry in data:
box = entry[0]
text = entry[1]
confidence = entry[2]
result = {
'box': box,
'text': text,
'confidence': confidence
}
results.append(result)
return results

def parse_paddle_result(data):
results = []
for entry in data[0]:
box = entry[0]
text = entry[1][0]
confidence = entry[1][1]
result = {
'box': box,
'text': text,
'confidence': confidence
}
results.append(result)
return results


106 changes: 85 additions & 21 deletions detectron/demo/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,21 @@
from gradio_image_prompter import ImagePrompter
from retriever import knowledgeBase
import sys
import os
sys.path.append("..")
from event_handler import components,gradio
from event_handler import create_event_handlers


from pathlib import Path

# 获取当前文件的绝对路径
current_file_path = Path(__file__).resolve()

# 获取当前文件所在的目录
current_directory = current_file_path.parent
print(current_directory)

def create_ui():
with gr.Blocks() as demo:
with gr.Tab("基础算法"):
Expand All @@ -25,11 +36,13 @@ def create_ui():
components["image_input"] = gr.Image(type="pil",elem_id='image-input',label='输入')
components["base_examples"] = gr.Examples(
examples=[
["./examples/horse.png"],
["./examples/view.jpeg"],
["./examples/cat.jpg"],
[os.path.join(current_directory,"examples/horse.png")],
[os.path.join(current_directory,"examples/view.jpeg")],
[os.path.join(current_directory,"examples/cat.jpg")],
],
inputs=components["image_input"],
examples_per_page=3,
label="图片示例"
)
with gr.Column(scale=2):
with gr.Row():
Expand Down Expand Up @@ -59,15 +72,24 @@ def create_ui():
with gr.Group():
components["yolo_image_input"] = gr.Image(type="pil",elem_id='image-input',label='输入')
components["yolo_video_input"] = gr.Video(label='输入',visible=False,interactive=True)
components["yolo_examples"] = gr.Examples(
components["yolo_image_examples"] = gr.Examples(
examples=[
[os.path.join(current_directory,"examples/horse.png")],
[os.path.join(current_directory,"examples/view.jpeg")],
[os.path.join(current_directory,"examples/cat.jpg")],
],
inputs=[components["yolo_image_input"]],
examples_per_page=3,
label="图片示例"
)
components["yolo_video_examples"] = gr.Examples(
examples=[
["./examples/horse.png",None],
["./examples/view.jpeg",None],
["./examples/cat.jpg",None],
[None,"./examples/gym.mp4"],
[None,"./examples/traffic.mp4"],
[os.path.join(current_directory,"examples/gym.mp4")],
[os.path.join(current_directory,"examples/traffic.mp4")],
],
inputs=[components["yolo_image_input"],components["yolo_image_input"]],
inputs=[components["yolo_video_input"]],
examples_per_page=2,
label="视频示例"
)

with gr.Column(scale=2):
Expand Down Expand Up @@ -97,10 +119,12 @@ def create_ui():
components["face_input"] = gr.Gallery(elem_id='face-input',label='输入',columns=2,type="pil")
components["face_examples"] = gr.Examples(
examples=[
["./examples/face1.jpeg"],
["./examples/face2.jpeg"],
[[(os.path.join(current_directory,"examples/face1.jpeg"),"face1")]],
[[(os.path.join(current_directory,"examples/face2.jpeg"),"face2")]],
],
inputs=components["face_input"],
inputs=[components["face_input"]],
examples_per_page=2,
label="图片示例"
)
with gr.Column(scale=2):
with gr.Row():
Expand All @@ -110,6 +134,37 @@ def create_ui():
with gr.Row():
with gr.Group():
components["face_output"] = gr.JSON(label="推理结果")
with gr.Tab("OCR"):
with gr.Row():
with gr.Column(scale=2):
components["ocr_type"] = gr.Dropdown(
["OCR","Easy"],value="Easy",
label="算法类别",interactive=True
)
with gr.Column(scale=2):
components["submit_ocr_btn"] = gr.Button(value="解析")

with gr.Row():
with gr.Column(scale=2):
with gr.Row(elem_id=''):
with gr.Group():
components["ocr_input"] = gr.Image(elem_id='ocr-input',label='输入',type="pil")
components["ocr_examples"] = gr.Examples(
examples=[
[os.path.join(current_directory,"examples/json.png")],
[os.path.join(current_directory,"examples/text.jpeg")],
],
inputs=[components["ocr_input"]],
examples_per_page=2,
label="图片示例"
)
with gr.Column(scale=2):
with gr.Row():
with gr.Group():
components["ocr_output"] = gr.Image(elem_id='ocr_output',label='输出',interactive=False,type="pil")
with gr.Row():
with gr.Group():
components["ocr_json_output"] = gr.JSON(label="推理结果")
with gr.Tab("SAM everything"):
with gr.Row():
with gr.Column(scale=2):
Expand All @@ -124,14 +179,23 @@ def create_ui():
with gr.Group():
components["sam_input"] = ImagePrompter(elem_id='sam-input',label='输入',type="pil")
components["sam_video_input"] = gr.Video(label='视频输入',visible=False,interactive=True)
components["face_examples"] = gr.Examples(
examples=[
["./examples/horse.png",None],
["./examples/view.jpeg",None],
[None,"./examples/traffic.mp4"],
],
inputs=[components["sam_input"],components["sam_video_input"]],
)
components["sam_image_examples"] = gr.Examples(
examples=[
[{'image': os.path.join(current_directory,"examples/horse.png"), 'points': []}],
[{'image': os.path.join(current_directory,"examples/view.jpeg"), 'points': []}],
],
inputs=[components["sam_input"]],
examples_per_page=2,
label="图片示例"
)
components["sam_video_examples"] = gr.Examples(
examples=[
[os.path.join(current_directory,"examples/traffic.mp4")],
],
inputs=[components["sam_video_input"]],
examples_per_page=1,
label="视频示例"
)
with gr.Column(scale=2):
with gr.Group():
components["sam_output"] = gr.Gallery(elem_id='sam_output',label='输出',columns=1,interactive=False)
Expand Down
4 changes: 3 additions & 1 deletion requirements-docker.txt
Original file line number Diff line number Diff line change
Expand Up @@ -23,4 +23,6 @@ shapely==2.0.5
langchain==0.2.12
langchain_community==0.2.11
transformers==4.44.0
lapx>=0.5.2
lapx>=0.5.2
paddleocr
easyocr
4 changes: 3 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -24,4 +24,6 @@ shapely==2.0.5
langchain==0.2.12
langchain_community==0.2.11
transformers==4.44.0
lapx>=0.5.2
lapx>=0.5.2
paddleocr
easyocr

0 comments on commit 2b270da

Please sign in to comment.