File size: 6,096 Bytes
38d4385
d9c0ea9
38d4385
 
890b83b
a9f1bab
 
 
 
1775e4d
d9c0ea9
38d4385
 
1775e4d
 
 
 
 
 
 
 
38d4385
a9f1bab
 
 
38d4385
 
 
 
890b83b
38d4385
fba92a9
765c987
890b83b
a9f1bab
38d4385
1f96cd2
765c987
1f96cd2
765c987
38d4385
a9f1bab
f769dab
 
a9f1bab
 
f769dab
a9f1bab
 
f769dab
a9f1bab
 
480872a
a9f1bab
480872a
f769dab
a9f1bab
f769dab
a9f1bab
 
 
 
 
 
 
 
 
 
 
 
f769dab
a9f1bab
5034720
a9f1bab
 
 
480872a
 
f769dab
 
 
 
5034720
 
 
 
 
a9f1bab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5034720
 
a9f1bab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d9a9d0e
a9f1bab
 
 
 
 
29a6070
a9f1bab
d9a9d0e
 
d9c0ea9
349b9d7
d9a9d0e
a9f1bab
d9c0ea9
 
a9f1bab
 
d9c0ea9
 
 
 
 
 
 
 
 
 
 
 
a9f1bab
 
890b83b
fba92a9
 
38d4385
890b83b
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
import tensorflow as tf
from PIL import Image
from io import BytesIO
from ultralytics import YOLO
import cv2
from datetime import datetime
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from pathlib import Path
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

labels = []
classification_model = tf.keras.models.load_model('./models.h5')
detection_model = YOLO('./best.pt')

with open("labels.txt") as f:
    for line in f:
        labels.append(line.replace('\n', ''))

def classify_image(img):
    # Resize the input image to the expected shape (224, 224)
    img_array = np.asarray(img.resize((224, 224)))[..., :3]
    img_array = img_array.reshape((1, 224, 224, 3))  # Add batch dimension
    img_array = tf.keras.applications.efficientnet.preprocess_input(img_array)
    prediction = classification_model.predict(img_array).flatten()
    confidences = {labels[i]: float(prediction[i]) for i in range(90)}
    # Sort the confidences dictionary by value and get the top 3 items
    # top_3_confidences = dict(sorted(confidences.items(), key=lambda item: item[1], reverse=True)[:3])
    
    return confidences

def animal_detect_and_classify(img_path):
    # Read the image using Pillow
    img = Image.open(img_path)

    # Pass the image through the detection model and get the result
    detect_results = detection_model(np.array(img))

    combined_results = []

    # Iterate over detections
    for result in detect_results:
        flag = False
        for box in result.boxes:
            flag = True
            # Crop the Region of Interest (RoI)
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            detect_img = img.crop((x1, y1, x2, y2)).resize((224, 224))

            # Convert the image to a numpy array
            inp_array = np.array(detect_img)

            # Reshape the array to match the expected input shape
            inp_array = inp_array.reshape((-1, 224, 224, 3))

            # Preprocess the input array
            inp_array = tf.keras.applications.efficientnet.preprocess_input(inp_array)

            # Make predictions using the classification model
            prediction = classification_model.predict(inp_array)

            # Map predictions to labels
            threshold = 0.66
            predicted_labels = [labels[np.argmax(pred)] if np.max(pred) >= threshold else "animal" for pred in prediction]
            print(predicted_labels)
            combined_results.append(((x1, y1, x2, y2), predicted_labels))
        if flag:
            continue
        # If no detections found, consider the whole image
        x2, y2 = img.size
        detect_img = img.resize((224, 224))
        inp_array = np.array(detect_img).reshape((-1, 224, 224, 3))
        inp_array = tf.keras.applications.efficientnet.preprocess_input(inp_array)
        prediction = classification_model.predict(inp_array)
        threshold = 0.66
        predicted_labels = [labels[np.argmax(pred)] if np.max(pred) >= threshold else "unknown" for pred in prediction]
        combined_results.append(((0, 0, x2, y2), predicted_labels))

    return combined_results

def generate_color(class_name):
    # Generate a hash from the class name
    color_hash = hash(class_name)
    print(color_hash)
    # Normalize the hash value to fit within the range of valid color values (0-255)
    color_hash = abs(color_hash) % 16777216
    R = color_hash//(256*256)
    G = (color_hash//256) % 256
    B = color_hash % 256
    # Convert the hash value to RGB color format
    color = (R, G, B)

    return color

def plot_detected_rectangles(image, detections, output_path):
    # Create a copy of the image to draw on
    img_with_rectangles = image.copy()

    # Iterate over each detected rectangle and its corresponding class name
    for rectangle, class_names in detections:
        if class_names[0] == "unknown":
            continue
        # Extract the coordinates of the rectangle
        x1, y1, x2, y2 = rectangle

        # Generate a random color
        color = generate_color(class_names[0])

        # Draw the rectangle on the image
        cv2.rectangle(img_with_rectangles, (x1, y1), (x2, y2), color, 2)

        # Put the class names above the rectangle
        for i, class_name in enumerate(class_names):
            cv2.putText(img_with_rectangles, class_name, (x1, y1 - 10 - i*20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    # Show the image with rectangles and class names
    cv2.imwrite(output_path, img_with_rectangles)


@app.post("/predict/v2")
async def predict_v2(file: UploadFile = File(...)):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_")
    filename = timestamp + file.filename
    contents = await file.read()
    image = Image.open(BytesIO(contents))
    image.save("input/" + filename)
    detections = animal_detect_and_classify("input/" + filename)
    class_names = [class_name[0] for _, class_name in detections]
    plot_detected_rectangles(cv2.imread("input/" + filename), detections, "output/" + filename)
    return {
        "message": "Detection and classification completed successfully",
        "out": filename,
        "class_names": class_names
    }

IMAGE_DIR = Path("output")

@app.get("/image/")
async def get_image(image_name: str):
    # Sanitize the image_name to prevent directory traversal attacks
    if "../" in image_name:
        raise HTTPException(status_code=400, detail="Invalid image name")

    # Construct the image path
    image_path = IMAGE_DIR / image_name

    # Check if the image exists
    if not image_path.exists() or not image_path.is_file():
        raise HTTPException(status_code=404, detail="Image not found")

    # Return the image file
    return FileResponse(image_path)

@app.post("/predict")
async def predict(file: bytes = File(...)):
    img = Image.open(BytesIO(file))
    confidences = classify_image(img)
    return confidences