Update app.py

app.py

@@ -25,49 +25,24 @@ id2label = {
 }
 
 def remove_background(image):
-    # Convert the image to RGBA
     image = image.convert("RGBA")
-    
-    # Remove the background
     image_np = np.array(image)
     output_np = rembg.remove(image_np)
-    
-    # Create a white background image
     white_bg = Image.new("RGBA", image.size, "WHITE")
-    
-    # Composite the original image over the white background
     output_image = Image.alpha_composite(white_bg, Image.fromarray(output_np))
-    
-    # Convert back to RGB
     output_image = output_image.convert("RGB")
-    
     return output_image
 
     
 def predict(image):
-    # Remove the background
     image = remove_background(image)
-
-    # Preprocess the image
     inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Forward pass through the model
     outputs = model(**inputs)
-    
-    # Get the logits
     logits = outputs.logits
-    
-    # Calculate confidence scores with softmax
     probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-    
-    # Get the index of the most probable class
     predicted_class_idx = probs.argmax().item()
-    
-    # Get the label and confidence score of the most probable class
     predicted_label = id2label[str(predicted_class_idx)]
-    confidence_score = probs[predicted_class_idx].item() * 100  # Multiply by 100 to get a percentage
-    
-    # Return the label and confidence score
+    confidence_score = probs[predicted_class_idx].item() * 100  
     return f"{predicted_label}: {confidence_score:.2f}%"
 
 # Create the Gradio interface