panda1835 commited on
Commit
39af238
·
1 Parent(s): 5c86e3d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +101 -0
app.py ADDED
@@ -0,0 +1,101 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from keras.models import load_model
2
+ from PIL import Image, ImageOps
3
+ import numpy as np
4
+ import gradio as gr
5
+ import pandas as pd
6
+ import torch
7
+ import json
8
+
9
+ # Load the model
10
+ classify_model = load_model('keras_model.h5')
11
+ detect_model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True, _verbose=False)
12
+
13
+ def detect(image):
14
+ # Inference
15
+ results = detect_model(image)
16
+ try:
17
+ results = json.loads(results.pandas().xyxy[0].to_json(orient="records"))[0]
18
+ top = int(results['ymin'])
19
+ left = int(results['xmin'])
20
+ width = int(results['xmax'] - results['xmin'])
21
+ height = int(results['ymax'] - results['ymin'])
22
+ return top, left, width, height
23
+ except:
24
+ return 0,0,0,0
25
+
26
+ def format_label(label):
27
+ """
28
+ From '0 class 1\n' to 'class 1'
29
+ """
30
+ return label[:-1]
31
+
32
+
33
+ def predict(image):
34
+ top, left, width, height = detect(image)
35
+
36
+ if (top == 0) and (left == 0) and (width == 0) and (height==0):
37
+ return {
38
+ "predictions": {},
39
+ 'bbox': {
40
+ "top": 0,
41
+ "left": 0,
42
+ "width": 0,
43
+ "height": 0
44
+ }
45
+ }
46
+
47
+ if width > height:
48
+ height = width
49
+ else:
50
+ width = height
51
+
52
+ # Crop the turtle
53
+ image = image.crop((left, top, left + width, top + height))
54
+
55
+ # Create the array of the right shape to feed into the keras model
56
+ # The 'length' or number of images you can put into the array is
57
+ # determined by the first position in the shape tuple, in this case 1.
58
+ data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)
59
+
60
+ #resize the image to a 224x224 with the same strategy as in TM2:
61
+ #resizing the image to be at least 224x224 and then cropping from the center
62
+ size = (224, 224)
63
+ image = ImageOps.fit(image, size, Image.ANTIALIAS)
64
+
65
+ #turn the image into a numpy array
66
+ image_array = np.asarray(image)
67
+ # Normalize the image
68
+ normalized_image_array = (image_array.astype(np.float32) / 127.0) - 1
69
+ # Load the image into the array
70
+ data[0] = normalized_image_array
71
+
72
+ # run the inference
73
+ pred = classify_model.predict(data)
74
+ pred = pred.tolist()
75
+
76
+ with open('labels.txt','r') as f:
77
+ labels = f.readlines()
78
+
79
+ result = {format_label(labels[i]): round(pred[0][i],2) for i in range(len(pred[0]))}
80
+ sorted_result = {k: v for k, v in sorted(result.items(), key=lambda item: item[1], reverse=True) if v > 0}
81
+
82
+
83
+ return json.dumps({
84
+ "predictions": sorted_result,
85
+ 'bbox': {
86
+ "top": top,
87
+ "left": left,
88
+ "width": width,
89
+ "height": height
90
+ }
91
+ })
92
+
93
+ title = "🐆"
94
+
95
+ gr.Interface(
96
+ fn=predict,
97
+ inputs=gr.Image(type="pil", label="Input Image"),
98
+ outputs=[gr.JSON()],
99
+ # live=True,
100
+ title=title,
101
+ ).launch(share=True, debug=False)