fix: bug

README.md

@@ -1,5 +1,5 @@
 ---
-title: Yolov8 Ndl Layout
+title: Yolov5 Ndl Layout
 emoji: 🐢
 colorFrom: indigo
 colorTo: red

app.py

@@ -1,27 +1,24 @@
 import gradio as gr
-import torch
 from PIL import Image
-import numpy as np
-
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
-
+import yolov5
 import json
 
-def yolo(im, size=1024):
-    g = (size / max(im.size))  # gain
-    im = im.resize((int(x * g) for x in im.size), Image.LANCZOS)  # resize
+model = yolov5.load("nakamura196/yolov5-ndl-layout")
+
+def yolo(im):
 
     results = model(im)  # inference
-    if results.imgs:
-        writable_img = np.copy(results.imgs[0])  # Create a writable copy of the image data
-        results.imgs[0] = writable_img
-        results.render()  # updates results.imgs with boxes and labels
 
     df = results.pandas().xyxy[0].to_json(orient="records")
     res = json.loads(df)
 
+    im_with_boxes = results.render()[0]  # results.render() returns a list of images
+    
+    # Convert the numpy array back to an image
+    output_image = Image.fromarray(im_with_boxes)
+
     return [
-        Image.fromarray(results.imgs[0]),
+        output_image,
         res
     ]
 
@@ -36,8 +33,6 @@ title = "YOLOv5 NDL-DocL Datasets"
 description = "YOLOv5 NDL-DocL Datasets Gradio demo for object detection. Upload an image or click an example image to use."
 article = "<p style='text-align: center'>YOLOv5 NDL-DocL Datasets is an object detection model trained on the <a href=\"https://github.com/ndl-lab/layout-dataset\">NDL-DocL Datasets</a>.</p>"
 
-# <a href=\"aaa\"></a>
-
 examples = [
     ['『源氏物語』(東京大学総合図書館所蔵).jpg'],
     ['『源氏物語』(京都大学所蔵).jpg'],

requirements.txt

@@ -1 +1 @@
-ultralyticsplus
+yolov5