MNGames commited on
Commit
99f44fe
1 Parent(s): af290e7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -65
app.py CHANGED
@@ -1,66 +1,24 @@
1
  import gradio as gr
2
- import torch
3
- import pytesseract
4
- from transformers import AutoTokenizer, AutoModelForSequenceClassification
5
-
6
- # Load the Hugging Face model for object detection
7
- model_name = "flax-community/yolov5s-v1-coco"
8
- model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
9
-
10
- # Load the Hugging Face model for text classification
11
- classification_model_name = "distilbert-base-uncased"
12
- classification_tokenizer = AutoTokenizer.from_pretrained(classification_model_name)
13
- classification_model = AutoModelForSequenceClassification.from_pretrained(classification_model_name)
14
-
15
- # Define function for text recognition (OCR)
16
- def perform_ocr(image):
17
- # Perform OCR using pytesseract
18
- text = pytesseract.image_to_string(image)
19
- return text
20
-
21
- # Define function to process video and predict
22
- def predict_runner_status(video_file):
23
- cap = cv2.VideoCapture(video_file.name)
24
-
25
- results = []
26
-
27
- while cap.isOpened():
28
- ret, frame = cap.read()
29
- if not ret:
30
- break
31
-
32
- # Object detection
33
- results_detection = model(frame)
34
-
35
- # Logic for determining runner status using detected objects
36
- objects = results_detection.pred[0][:, -1].numpy()
37
- if 0 in objects: # 0 corresponds to person class
38
- # Perform OCR on the detected person
39
- person_bbox = results_detection.pred[0][np.where(objects == 0)][0][:4]
40
- person_bbox = person_bbox.astype(int)
41
- person_img = frame[person_bbox[1]:person_bbox[3], person_bbox[0]:person_bbox[2]]
42
- text = perform_ocr(person_img)
43
-
44
- # Classification using text classification model
45
- inputs_classification = classification_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
46
- outputs_classification = classification_model(**inputs_classification)
47
- predicted_class = torch.argmax(outputs_classification.logits).item()
48
- if predicted_class == 1:
49
- runner_status = "Out"
50
- else:
51
- runner_status = "Safe"
52
-
53
- result = {
54
- "frame_number": cap.get(cv2.CAP_PROP_POS_FRAMES),
55
- "runner_status": runner_status
56
- }
57
- results.append(result)
58
-
59
- cap.release()
60
-
61
- return results
62
-
63
- inputs = gr.inputs.Video(type="file", label="Upload a baseball video")
64
- outputs = gr.outputs.Label(type="auto", label="Runner Status")
65
- interface = gr.Interface(fn=predict_runner_status, inputs=inputs, outputs=outputs, title="Baseball Runner Status Predictor")
66
- interface.launch(share=True)
 
1
  import gradio as gr
2
+ from transformers import pipeline
3
+
4
+ # Replace with a suitable image classification model ID
5
+ model_id = "sayakpaul/resnet-50-finetuned-imagenet"
6
+
7
+ def analyze_image(image):
8
+ classifier = pipeline("image-classification", model=model_id)
9
+ predictions = classifier(images=image) # Assuming the model outputs probabilities
10
+ # Extract the most likely class and its probability
11
+ top_class = predictions[0]["label"]
12
+ top_prob = predictions[0]["score"]
13
+ return f"Top Class: {top_class} (Probability: {top_prob:.2f})"
14
+
15
+ # Gradio interface
16
+ interface = gr.Interface(
17
+ fn=analyze_image,
18
+ inputs="image",
19
+ outputs="text",
20
+ title="Image Analyzer (Generic)",
21
+ description="Upload an image and get the most likely classification based on the chosen model.",
22
+ )
23
+
24
+ interface.launch()