diff --git a/Backend/backend.py b/Backend/backend.py index 5d29986..8891ee2 100644 --- a/Backend/backend.py +++ b/Backend/backend.py @@ -10,7 +10,7 @@ import requests from flask_cors import CORS, cross_origin -from mongoengine import connect, Document, StringField, DoesNotExist +from mongoengine import connect, Document, StringField, BinaryField, DoesNotExist from werkzeug.security import generate_password_hash, check_password_hash import secrets # For generating a session key @@ -20,6 +20,8 @@ from db_access import Image from db_access import check_user, create_user, db_connect, get_image, create_image +import base64 + app = Flask(__name__) cors = CORS(app) @@ -34,10 +36,7 @@ @app.route("/generate_image", methods=["POST"]) -@jwt_required() def generate_image(): - current_user = get_jwt_identity() - data = request.get_json() url = "https://imagegolf.io/api/generate" url_data = {"inputValue": data["prompt"]} @@ -85,6 +84,7 @@ def update_image_elo(): @jwt_required() def store_image(): current_user = get_jwt_identity() # Gets the identity of the current user + print(current_user) data = request.get_json() # Validate required fields @@ -93,11 +93,13 @@ def store_image(): if field not in data: return jsonify({"error": f"Missing field: {field}"}), 400 + image_res = requests.get(data["url"]) + # Prepare the data for creating an image image_data = { "creator": current_user, # Assuming the creator is the logged-in user "prompt": data["prompt"], - "url": data["url"] + "data": image_res.content } # Call the create_image function from db_access @@ -115,22 +117,20 @@ def store_image(): @app.route("/fetch_portfolio", methods=["GET"]) @jwt_required() def fetch_portfolio(): - current_user_username = get_jwt_identity() # Get the username from the JWT token - try: # Retrieve the user by username - user = User.objects.get(username=current_user_username) + user = get_jwt_identity() # Fetch all images created by this user - user_images = Image.objects(creator=user) + portfolio_images = User.objects.get(pk=user).portfolio # Format the response with the list of images portfolio = [{ "image_id": str(image.id), "prompt": image.prompt, - "url": image.url, - "creator": user - } for image in user_images] + "creator": user, + "data": base64.b64encode(image.data).decode("ascii") + } for image in portfolio_images] return jsonify(portfolio), 200 @@ -154,7 +154,7 @@ def login(): # Generate session key/token # This is just a placeholder for an actual session key/token # session_key = secrets.token_hex(16) - access_token = create_access_token(identity=str(user.username)) + access_token = create_access_token(identity=str(user.pk)) # You would store this session key in a session store or database # with a reference to the user and a valid time period @@ -177,7 +177,6 @@ def login(): ) except DoesNotExist: # Username does not exist - print("bad user") return ( jsonify({"message": "Login failed, invalid username"}), 401, diff --git a/Backend/db_access.py b/Backend/db_access.py index b994ebc..b5a0289 100644 --- a/Backend/db_access.py +++ b/Backend/db_access.py @@ -4,6 +4,7 @@ connect, Document, StringField, + BinaryField, IntField, ListField, ReferenceField, @@ -47,8 +48,8 @@ class User(Document): class Image(Document): creator = ReferenceField(User, required=True) - url = StringField(required=True) prompt = StringField(required=True) + data = BinaryField(required=True) meta = {"collection": "images"} @@ -60,7 +61,6 @@ def create_user(data): encrypted_password=data["password"], email=data["email"], portfolio=[] # Initialize an empty portfolio - ) try: user.save() @@ -91,16 +91,15 @@ def check_user(data): def create_image(data): try: - # Retrieve the user by username - creator = User.objects.get(username=data["creator"]) - # Create and save the image image = Image( - creator=creator, + creator=data["creator"], prompt=data["prompt"], - url=data["url"] + data=data["data"] ) image.save() + + creator = User.objects.get(pk=data["creator"]) # Optionally, you can also append this image to the user's portfolio here creator.update(push__portfolio=image) @@ -109,7 +108,7 @@ def create_image(data): except DoesNotExist: return Response("Creator user does not exist.", 404) except Exception as e: - return Response(f"Internal server error {e}", 500) + return Response(f"Internal server error: {e}", 500) def get_image(image_id): diff --git a/Frontend/src/pages/create.js b/Frontend/src/pages/create.js index 7aa006c..6b79b2e 100644 --- a/Frontend/src/pages/create.js +++ b/Frontend/src/pages/create.js @@ -1,36 +1,53 @@ import React, { useState } from "react"; import axios from "axios"; import Image from "next/image"; +import Cookie from "js-cookie"; -export default function Home() { +export default function Create() { const [text, setText] = useState(""); const [url, setUrl] = useState(""); const [isImageAccepted, setImageAccepted] = useState(null); // null = not decided, true = accepted, false = rejected + const [generating, setGenerating] = useState(false); + const [failed, setFailed] = useState(false); const handleSubmit = async (e) => { e.preventDefault(); try { - const response = await axios.post( + setGenerating(true); + setFailed(false); + await axios.post( "http://localhost:5000/generate_image", { prompt: text }, - ); + ).then(response => { + if (response.data.output) { + setUrl(response.data.output); + setImageAccepted(null); // Reset the decision state when a new image is fetched + } else if (response.data.error) { + setFailed(true); + console.error(response.data.error); + } + }).then(() => setGenerating(false)); - if (response.data.output) { - setUrl(response.data.output); - setImageAccepted(null); // Reset the decision state when a new image is fetched - } else if (response.data.error) { - console.error(response.data.error); - } } catch (error) { console.error("Error:", error); } }; - const handleAccept = () => { + const handleAccept = async (e) => { + e.preventDefault(); + setImageAccepted(true); - // You can add any additional logic here for when the image is accepted - }; + await axios.post("http://localhost:5000/store_image", + { + "prompt": text, + "url": url + }, + { + headers: { + "Authorization": `Bearer ${Cookie.get("token")}` + } + }).catch((error) => console.error("Error: ", error))}; const handleReject = () => { setImageAccepted(false); @@ -52,9 +69,10 @@ export default function Home() { /> @@ -90,6 +108,9 @@ export default function Home() { {isImageAccepted === false && (
Image rejected.
)} + {failed === true && ( +Image generation failed.
+ )} ); } diff --git a/Frontend/src/pages/nav.js b/Frontend/src/pages/nav.js index 105ae99..3e7e612 100644 --- a/Frontend/src/pages/nav.js +++ b/Frontend/src/pages/nav.js @@ -19,12 +19,6 @@ const NavBar = () => { {/* Primary Navbar items */}