Update app.py

app.py

@@ -1,84 +1,53 @@
+import gradio as gr
 import torch
-import json
 from PIL import Image
-from pathlib import Path
-import gradio as gr
-import numpy as np
-import cv2
+import json
 from ultralytics import YOLO
 
-# 載入車框權重
-car_mask_weights = 'mask.pt'
-car_mask_model = torch.hub.load('./yolov5', 'custom', path=car_mask_weights, source="local")
-car_mask_model.conf = 0.7  # 車框信心指數
-
-# 載入車號權重
-plate_weights = 'plate.pt'
-plate_model = torch.hub.load('./yolov5', 'custom', path=plate_weights, source="local")
-plate_model.conf = 0.5  # 車號信心指數
-plate_model.classes = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35]  # 過濾'-'符號
-
-# 設定輸入和輸出資料夾
-input_folder = 'input/'
-output_folder = 'output/'
-
-# 確保輸出資料夾存在
-Path(output_folder).mkdir(parents=True, exist_ok=True)
-
-# 將預測邏輯封裝成函數
-def predict_license_plate(im):
-    # 使用車框權重進行預測
-    car_mask_results = car_mask_model(im, size=1280)
-
-    # 取得車框預測的座標和類別
-    car_mask_boxes = car_mask_results.pandas().xyxy[0]
-
-    # 對每個車框進行車號預測
-    plate_json = {}  # 將車號依序存成json
-    for i, box in car_mask_boxes.iterrows():
-        # 切割車框圖片
-        x1, y1, x2, y2 = box['xmin'], box['ymin'], box['xmax'], box['ymax']
-
-        # 放寬車牌裁切範圍
-        expand_factor = 0.1  # 調整擴大比例
-        expand_width = int((x2 - x1) * expand_factor)
-        expand_height = int((y2 - y1) * expand_factor)
-        x1 -= expand_width
-        y1 -= expand_height
-        x2 += expand_width
-        y2 += expand_height
-
-        # 限制裁切範圍在圖片邊界內
-        x1 = max(x1, 0)
-        y1 = max(y1, 0)
-        x2 = min(x2, im.width)
-        y2 = min(y2, im.height)
-
-        car_mask_image = im.crop((x1, y1, x2, y2))
-
-        # 使用車號權重進行預測
-        plate_results = plate_model(car_mask_image, size=640)
-
-        # 提取車號預測的結果
-        plate_labels = plate_results.pandas().xyxy[0].sort_values('xmin')['name'].tolist()
-
-        # 儲存圖片和預測結果
-        plate_image_name = f'{Path(image_file).stem}_plate{i}.jpg'
-        plate_image_path = Path(output_folder) / plate_image_name
-        car_mask_image.save(plate_image_path)  # 儲存圖片
-
-        licence_plate = ''.join(plate_labels)
-        plate_json[plate_image_name] = licence_plate
-
-    return [Image.fromarray(car_mask_results.ims[0]), plate_json]
-
-# 將函數包裝成 Gradio 介面
-inputs = gr.Image(type='pil', label="Original Image")
-outputs = [gr.Image(type="pil", label="Output Image"),
-           gr.JSON(label="Output JSON")]
-
-title = "License_Plate_Prediction"
-description = "Predict license plates using YOLOv5."
-examples = [Image.open('1.jpg'), Image.open('2.jpg')]
-
-gr.Interface(predict_license_plate, inputs, outputs, title=title, description=description, examples=examples).launch(enable_queue=True)
+
+# Images
+torch.hub.download_url_to_file(
+    'https://i.imgur.com/4GmZXID.jpg', '1.jpg')
+torch.hub.download_url_to_file(
+    'https://i.imgur.com/ktIGRvs.jpg', '2.jpg')
+torch.hub.download_url_to_file(
+    'https://i.imgur.com/fSEsXoE.jpg', '3.jpg')
+torch.hub.download_url_to_file(
+    'https://i.imgur.com/lsVJRzd.jpg', '4.jpg')
+torch.hub.download_url_to_file(
+    'https://i.imgur.com/1OFmJd1.jpg', '5.jpg')
+torch.hub.download_url_to_file(
+    'https://i.imgur.com/GhfAWMJ.jpg', '6.jpg')
+
+# Model
+# model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # force_reload=True to update
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='plate.pt', source="local") 
+
+
+def yolo(im):
+    model.conf = 0.6  # NMS confidence threshold
+    # g = (size / max(im.size))  # gain
+    # im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize
+
+    results = model(im, size=1280)  # inference
+    results.render()  # updates results.imgs with boxes and labels
+
+    df = results.pandas().xyxy[0].sort_values('xmin')[['name']].to_json(orient="records") # 可以把[['name']]刪除即可顯示全部
+    res = json.loads(df)
+
+    return [Image.fromarray(results.ims[0]), res]
+
+
+
+
+
+inputs = gr.inputs.Image(type='pil', label="Original Image")
+outputs = [gr.outputs.Image(type="pil", label="Output Image"),
+    gr.outputs.JSON(label="Output JSON")]
+
+title = "TW_plate_number"
+description = "TW_plate_number"
+
+
+examples = [['1.jpg'], ['2.jpg'], ['3.jpg'], ['4.jpg'], ['5.jpg'], ['6.jpg']]
+gr.Interface(yolo, inputs, outputs, title=title, description=description, examples=examples, theme="huggingface").launch(enable_queue=True)