Spaces:
Sleeping
Sleeping
File size: 3,008 Bytes
e785be9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import os
import sys
import uuid
# Fix path
sys.path.append(os.getcwd())
from app.services.memory_service import memory_service
from app.services.object_service import detector
from tqdm import tqdm
BASE_DIR = r"c:\Users\Krish\Downloads\Convolve\MYNursingHome"
LOCATION_MAP = {
"bed": "Bedroom 101",
"chair": "Common Area",
"medicine_box": "Nurse Station",
"keys": "Reception Desk",
"wheelchair": "Entrance Lobby",
"walker": "Corridor A",
"clock": "Wall (Hallway)",
"phone": "Living Room",
"spectacles": "Bedside Table",
"remote": "TV Room",
"water_bottle": "Kitchen",
"shoes": "Shoe Rack",
"wallet": "Safe Box",
"book": "Library Shelf",
"basket_bin": "Corner",
"bench": "Garden",
"cabinet": "Storage Room",
"call_bell": "Bedside",
"cane_stick": "Entrance",
"door": "Main Entrance",
"electric_socket": "Wall",
"fan": "Ceiling",
"fire_extinguisher": "Hallway B",
"handrail": "Stairs",
"human_being": "Everywhere",
"rack": "Store",
"refrigerator": "Kitchen",
"shower": "Bathroom",
"sink": "Washroom",
"sofa": "Lounge",
"table": "Dining Hall",
"television": "TV Room",
"toilet_seat": "Restroom",
"wardrobe": "Bedroom 101",
"water_dispencer": "Corridor B"
}
def train_objects():
print("--- Training Object Memory ---")
# Iterate Categories
if not os.path.exists(BASE_DIR):
print(f"Directory not found: {BASE_DIR}")
return
categories = [d for d in os.listdir(BASE_DIR) if os.path.isdir(os.path.join(BASE_DIR, d))]
total_stored = 0
for cat in tqdm(categories, desc="Categories"):
cat_path = os.path.join(BASE_DIR, cat)
location = LOCATION_MAP.get(cat.lower(), "General Storage")
# Iterate Images
images = [f for f in os.listdir(cat_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
# Limit to 5 images per category to save time/space for prototype
# Or do all? Let's do 10.
for img_name in images[:10]:
img_path = os.path.join(cat_path, img_name)
try:
# Generate Embedding
embedding = detector.generate_embedding(img_path)
# Store
metadata = {
"name": cat,
"type": "object",
"location": location,
"filename": img_name
}
memory_service.store_object_memory(
object_id=str(uuid.uuid4()),
embedding=embedding,
metadata=metadata
)
total_stored += 1
except Exception as e:
print(f"Error processing {img_name}: {e}")
print(f"--- Training Complete. Stored {total_stored} objects. ---")
if __name__ == "__main__":
train_objects()
|