Updated app.py

app.py

@@ -2,36 +2,143 @@
 
 import gradio as gr
 import tensorflow as tf
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
 import os
 
 # Suppress TensorFlow logging for cleaner logs
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
-def list_tensor_names():
-    # Load the frozen TensorFlow graph
-    MODEL_DIR = 'tensorflow/'
-    MODEL_PATH = os.path.join(MODEL_DIR, 'model.pb')
-    
-    with tf.io.gfile.GFile(MODEL_PATH, 'rb') as f:
+# Load labels
+with open('tensorflow/labels.txt', 'r') as f:
+    labels = f.read().splitlines()
+
+# Function to load the frozen TensorFlow graph
+def load_frozen_graph(pb_file_path):
+    with tf.io.gfile.GFile(pb_file_path, 'rb') as f:
         graph_def = tf.compat.v1.GraphDef()
         graph_def.ParseFromString(f.read())
     
     with tf.Graph().as_default() as graph:
         tf.import_graph_def(graph_def, name='')
+    return graph
+
+# Load the TensorFlow model
+MODEL_DIR = 'tensorflow/'
+MODEL_PATH = os.path.join(MODEL_DIR, 'model.pb')
+graph = load_frozen_graph(MODEL_PATH)
+sess = tf.compat.v1.Session(graph=graph)
+
+# Define tensor names based on your model's outputs
+input_tensor = graph.get_tensor_by_name('image_tensor:0')
+detected_boxes = graph.get_tensor_by_name('detected_boxes:0')
+detected_classes = graph.get_tensor_by_name('detected_classes:0')
+detected_scores = graph.get_tensor_by_name('detected_scores:0')
+
+# Define the target size based on your model's expected input
+# You may need to adjust these values based on your model's requirements
+TARGET_WIDTH = 224
+TARGET_HEIGHT = 224
+
+def preprocess_image(image):
+    """
+    Preprocess the input image:
+    - Resize to target dimensions
+    - Convert to numpy array
+    - Normalize pixel values
+    - Convert RGB to BGR if required by the model
+    """
+    image = image.resize((TARGET_WIDTH, TARGET_HEIGHT))
+    image_np = np.array(image).astype(np.float32)
+    image_np = image_np / 255.0  # Normalize to [0,1]
+    image_np = image_np[:, :, (2, 1, 0)]  # Convert RGB to BGR if required
+    image_np = np.expand_dims(image_np, axis=0)  # Add batch dimension
+    return image_np
+
+def draw_boxes(image, boxes, classes, scores, threshold=0.5):
+    """
+    Draw bounding boxes and labels on the image.
+    Args:
+        image (PIL.Image): The original image.
+        boxes (np.array): Array of bounding boxes.
+        classes (np.array): Array of class IDs.
+        scores (np.array): Array of confidence scores.
+        threshold (float): Confidence threshold to filter detections.
+    Returns:
+        PIL.Image: Annotated image.
+    """
+    draw = ImageDraw.Draw(image)
+    try:
+        font = ImageFont.truetype("arial.ttf", 15)
+    except IOError:
+        font = ImageFont.load_default()
+    
+    for box, cls, score in zip(boxes[0], classes[0], scores[0]):
+        if score < threshold:
+            continue
+        # Convert box coordinates from normalized to absolute
+        ymin, xmin, ymax, xmax = box
+        left = xmin * image.width
+        right = xmax * image.width
+        top = ymin * image.height
+        bottom = ymax * image.height
+        
+        # Draw rectangle
+        draw.rectangle([(left, top), (right, bottom)], outline="red", width=2)
+        
+        # Prepare label
+        label = f"{labels[int(cls) - 1]}: {score:.2f}"
+        
+        # Draw label background
+        text_size = draw.textsize(label, font=font)
+        draw.rectangle([(left, top - text_size[1]), (left + text_size[0], top)], fill="red")
+        
+        # Draw text
+        draw.text((left, top - text_size[1]), label, fill="white", font=font)
+    
+    return image
+
+def predict(image):
+    """
+    Perform inference on the input image and return the annotated image.
+    Args:
+        image (PIL.Image): Uploaded image.
+    Returns:
+        PIL.Image: Annotated image with bounding boxes and labels.
+    """
+    try:
+        # Preprocess the image
+        input_array = preprocess_image(image)
+        
+        # Run inference
+        boxes, classes, scores = sess.run(
+            [detected_boxes, detected_classes, detected_scores],
+            feed_dict={input_tensor: input_array}
+        )
+        
+        # Annotate the image with bounding boxes and labels
+        annotated_image = draw_boxes(image.copy(), boxes, classes, scores, threshold=0.5)
+        
+        return annotated_image
     
-    tensor_names = [op.name for op in graph.get_operations()]
-    return "\n".join(tensor_names)
+    except Exception as e:
+        return Image.new('RGB', (500, 500), color = (255, 0, 0)).convert('RGB').copy()
 
 # Define Gradio interface
-title = "JunkWaxHero 🦸‍♂️ - Tensor Names Inspector"
-description = "This interface lists all tensor names in the TensorFlow model to help identify the correct input and output tensors."
+title = "JunkWaxHero 🦸‍♂️ - Baseball Card Set Identifier"
+description = "Upload an image of a baseball card, and JunkWaxHero will identify the set it belongs to with high accuracy."
 
 iface = gr.Interface(
-    fn=list_tensor_names,
-    inputs=None,
-    outputs="text",
+    fn=predict,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=gr.outputs.Image(type="pil"),
     title=title,
     description=description,
+    examples=[
+        ["examples/card1.jpg"],
+        ["examples/card2.jpg"],
+        ["examples/card3.jpg"]
+    ],
     allow_flagging="never"
 )