KristofGaming39's picture
Update app.py
6afcf4a verified
import gradio as gr
import requests
from bs4 import BeautifulSoup
from PIL import Image
from transformers import ViltProcessor, ViltForQuestionAnswering
import torch
import torch.nn.functional as F
# ViLT Model betöltése
model_path = "dandelin/vilt-b32-finetuned-vqa" # Hugging Face modell
device = "cuda" if torch.cuda.is_available() else "cpu"
model = ViltForQuestionAnswering.from_pretrained(model_path).to(device)
processor = ViltProcessor.from_pretrained(model_path)
# Keresés a Roblox katalógusban
def search_roblox(character_name):
url = f"https://www.roblox.com/catalog?Category=3&salesTypeFilter=1"
headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36'}
response = requests.get(url, headers=headers)
if response.status_code != 200:
return None, "Failed to fetch data", 0
soup = BeautifulSoup(response.text, "html.parser")
# Az elemzéshez szükséges adatokat a megfelelő HTML elemekből szűrjük
items = soup.find_all('div', {'class': 'catalog-item-container'})
best_item = None
best_score = 0
for item in items[:5]: # Csak az első 5 találatot nézzük
try:
img_tag = item.find('img')
img_url = img_tag['src'] if img_tag else None
name = item.find('div', {'class': 'item-card-name'}).text.strip()
link_tag = item.find('a', {'class': 'item-card-link'})
link = link_tag['href'] if link_tag else None
if not img_url or not link:
continue
image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
# AI kiértékelés
question = f"Is this item good for the character: {character_name}? Give a score from 0 to 1000."
inputs = processor(images=image, text=question, return_tensors="pt").to(device)
outputs = model(**inputs)
softmax_scores = F.softmax(outputs.logits, dim=-1)
score = softmax_scores.max().item() * 1000 # 0-1000 közötti skála
if score > best_score:
best_score = score
best_item = {"name": name, "img_url": img_url, "link": link, "score": best_score}
except Exception as e:
print(f"Error processing item: {e}")
if best_item:
return best_item["img_url"], best_item["link"], best_item["score"]
else:
return None, "No items found", 0
# Gradio UI
iface = gr.Interface(
fn=search_roblox,
inputs=gr.Textbox(label="Enter Character Name"),
outputs=[gr.Image(label="Best Match"), gr.Textbox(label="Item Link"), gr.Number(label="AI Score")],
title="Roblox Character Item Finder",
description="Enter a character name, and the AI will find the best Roblox catalog item for it."
)
iface.launch()