Update app.py

app.py

@@ -1,55 +1,83 @@
-import gradio as gr
 import torch
-from PIL import Image
 import json
-from ultralytics import YOLO
-
-
-# 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"
-# article = "<p style='text-align: center'>TW_plate_number <a href=\"http://codh.rois.ac.jp/char-shape/\">日本古典籍くずし字データセット</a>.</p>"
-
-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) # cache_examples=True,
-
-
+from PIL import Image
+from pathlib import Path
+import gradio as gr
+import numpy as np
+import cv2
+
+# 載入車框權重
+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('input/1.jpg'), Image.open('input/2.jpg')]
+
+gr.Interface(predict_license_plate, inputs, outputs, title=title, description=description, examples=examples).launch(enable_queue=True)