Adds more subcategories;

sepinf-inc · Feb 24, 2025 · 676837f · 676837f
1 parent 9d918a4
commit 676837f
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 11 deletions.
diff --git a/iped-app/resources/config/conf/CategoriesConfig.json b/iped-app/resources/config/conf/CategoriesConfig.json
@@ -202,10 +202,17 @@
 	{"name": "Classified by AI", "categories":[
       {"name": "AI Label: Child Sexual Abuse"},
       {"name": "AI Label: Likely Child Sexual Abuse"},
-      {"name": "AI Label: Drawing"},
+      {"name": "AI Label: Drawing", "categories":[
+      	{"name":"AI Label: Child Sexual Abuse Drawing"},
+      	{"name":"AI Label: AI Label: Explicit Drawing"}
+      ]},
       {"name": "AI Label: People"},
       {"name": "AI Label: Pornography"},
-      {"name": "AI Label: Other"}
+      {"name": "AI Label: Other", "categories":[
+      	{"name":"AI Label: Banknotes"},
+      	{"name":"AI Label: Digitalization"},
+      	{"name":"AI Label: Documents"}
+      ]}
     ]}
   ]
 }
diff --git a/iped-app/resources/scripts/tasks/AIRefineCategoryTask.js b/iped-app/resources/scripts/tasks/AIRefineCategoryTask.js
@@ -30,20 +30,45 @@ function finish(){}
  * Changes category of items based on their properties
  *
  */
+
+
+function addCategories(e, listcategories, default_cat) {
+    default_cat = (typeof default_cat !== "undefined") ? default_cat : null;
+	var tot=0;
+	for(var cat in listcategories){
+		if (e.getExtraAttribute(cat)>categorizationThreshold){
+			if(typeof listcategories[cat] ==='string'){
+				e.addCategory(listcategories[cat]);	 // Add category if it's a string
+				tot++;
+			}else{
+				tot+=addCategories(e,listcategories[cat]["subCategories"],listcategories[cat]["name"]);
+			}
+		}
+	}
+	if(tot===0 && default_cat !== null){
+		e.addCategory(default_cat);
+	}
+}
+
 function process(e){
 
-	listcategories={
+	var listcategories={
 		"AI_CSAM":"AI Label: Child Sexual Abuse",
 		"AI_LIKELYCSAM":"AI Label: Likely Child Sexual Abuse",
-		"AI_Drawing":"AI Label: Drawing",
-		"AI_Other":"AI Label: Other",
+		"AI_Drawing":{"name":"AI Label: Drawing","subCategories":{
+			"AI_Drawing_CSAM":"AI Label: Child Sexual Abuse Drawing",
+			"AI_Drawing_Porn":"AI Label: Explicit Drawing"}
+		},
+		"AI_Other":{"name":"AI Label: Other","subCategories":{
+			"AI_Other_Banknotes":"AI Label: Banknotes",
+			"AI_Other_Digitalization":"AI Label: Digitalization",
+			"AI_Other_Documents":"AI Label: Documents"}
+		},
 		"AI_People":"AI Label: People",
-		"AI_Porn":"AI Label: Pornography",
+		"AI_Porn":"AI Label: Pornography"
+
 		}
-	for(cat in listcategories){
-		if (e.getExtraAttribute(cat)>categorizationThreshold){
-			e.addCategory(listcategories[cat]);
-		}		
-	}
+
+	addCategories(e,listcategories);
 }