Photon08's picture
Update app.py
6d604b5
from ultralytics import YOLO
import random
import streamlit as st
from PIL import Image
import os
import gdown
def extract_user_input(image=None):
id = "1e2pHejcp96q-NT9VmIDEY_hCA-tiFNef"
output = "trained_model.pt"
gdown.download(id = id, output = output, quiet=False)
model = YOLO(output)
result = model.predict(image)
names = model.names
det_cls = []
for r in result:
for c in r.boxes.cls:
det_cls.append(names[int(c)])
return det_cls[0]
def ai_choice():
choice = random.randint(0,1000000000)
if choice <=100000:
choice = "Rock"
elif choice >101 and choice <=1000000:
choice = "Paper"
else:
choice = "Scissors"
return choice
def winner(user_input,ai_input):
if user_input == "Rock":
if ai_input == "Scissors":
return "Player 1"
elif ai_input == "Paper":
return "AI"
else:
return "Draw!!!"
if user_input == "Paper":
if ai_input == "Rock":
return "Player 1"
elif ai_input == "Scissors":
return "AI"
else:
return "Draw!!!"
if user_input == "Scissors":
if ai_input == "Rock":
return "AI"
elif ai_input == "Paper":
return "Player 1"
else:
return "Draw!!!"
def main():
col1, col2 = st.columns(2)
with st.container():
st.header("Welcome to the Computer Vison powered Rock-Paper-Scissors Game!")
st.caption("developed by Indranil Bhattacharyya")
player_name = st.text_input("Please enter your name, Player!: ")
image = st.camera_input("Please upload your image: ")
if image:
st.image(image)
im = Image.open(image)
rgb_im = im.convert("RGB")
# exporting the image
rgb_im.save("user_input.jpg")
user_input = extract_user_input("user_input.jpg")
user_input_front_end = player_name + " has selected: " + user_input
st.code(user_input_front_end)
ai_choice_var = ai_choice()
ai_choice_front_end = "AI has selected: " + ai_choice_var
st.code(ai_choice_front_end)
if len(user_input)>0:
winner_name = winner(user_input,ai_choice_var)
if len(winner_name)!=0:
with st.container():
if winner_name == "Player 1":
winner_front_end = player_name + " Wins!!!"
st.success(winner_front_end)
else:
if winner_name == "AI":
winner_front_end = winner_name + " Wins!!!"
else:
winner_front_end = "It's a " + winner_name
st.error(winner_front_end)
else:
winner_name = ''
if __name__ == "__main__":
main()