LabelMappingLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.db.image;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers
{
    public class LabelMappingLayer<T> : NeuronLayer<T>
    {
        IXImageDatabaseBase m_db;
        string m_strSource = null;
        int m_nSourceId = 0;
        int m_nProjectID = 0;
        Dictionary<int, int> m_rgActualMappedLabelCounts = new Dictionary<int, int>();
        object m_syncActualMappedLabels = new object();
        int m_nLabelCount = 0;
 
        public LabelMappingLayer(CudaDnn<T> cuda, Log log, LayerParameter p, IXDatabaseBase db)
            : base(cuda, log, p)
        {
            if (db.GetVersion() != DB_VERSION.IMG_V1 && db.GetVersion() != DB_VERSION.IMG_V2)
                throw new Exception("Currently only image databases are supported by the LabelMappingLayer.");
 
            m_db = (IXImageDatabaseBase)db;
            m_type = LayerParameter.LayerType.LABELMAPPING;
        }
 
        public override void SetNetParameterUsed(NetParameter np)
        {
            base.SetNetParameterUsed(np);
 
            m_strSource = null;
            m_nSourceId = 0;
            m_nProjectID = np.ProjectID;
 
            foreach (LayerParameter p in np.layer)
            {
                if (p.type == LayerParameter.LayerType.DATA)
                {
                    m_strSource = p.data_param.source;
                    break;
                }
            }
 
            if (m_strSource != null)
                m_nSourceId = m_db.GetSourceID(m_strSource);
        }
 
        public string GetActualLabelCounts(string strSrc)
        {
            if (m_nLabelCount == 0)
                m_nLabelCount = m_db.GetLabels(m_nSourceId).Count;
 
            List<KeyValuePair<int, int>> rgKv;
 
            lock (m_syncActualMappedLabels)
            {
                rgKv = m_rgActualMappedLabelCounts.OrderBy(p => p.Key).ToList();
            }
 
            string str = "";
            int nIdx = 0;
 
            for (int i = 0; i < m_nLabelCount; i++)
            {
                if (nIdx < rgKv.Count)
                {
                    str += rgKv[nIdx].Key.ToString() + "->" + rgKv[nIdx].Value.ToString();
                    nIdx++;
                }
                else
                {
                    str += "0";
                }
 
                str += ", ";
            }
 
            return str.TrimEnd(',', ' ');
        }
 
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.LayerSetUp(colBottom, colTop);
 
            if (m_strSource != null && m_nProjectID > 0)
            {
                int nSrcId = m_nSourceId;
 
                if (nSrcId > 0)
                {
                    Dictionary<int, int> rgMapCounts = new Dictionary<int, int>();
                    Dictionary<int, List<int>> rgMaps = new Dictionary<int, List<int>>();
                    bool bHasConditions = false;
 
                    foreach (LabelMapping map in m_param.labelmapping_param.mapping)
                    {
                        if (!rgMapCounts.Keys.Contains(map.NewLabel))
                            rgMapCounts.Add(map.NewLabel, 2);
                        else
                            rgMapCounts[map.NewLabel]++;
 
                        if (!rgMaps.Keys.Contains(map.NewLabel))
                            rgMaps.Add(map.NewLabel, new List<int>());
 
                        rgMaps[map.NewLabel].Add(map.OriginalLabel);
 
                        if (map.ConditionBoostEquals.HasValue)
                            bHasConditions = true;
                    }
 
                    if (m_param.labelmapping_param.update_database)
                    {
                        string strLabelCounts = m_db.GetLabelCountsAsTextFromSourceId(nSrcId);
 
                        if (m_param.labelmapping_param.reset_database_labels)
                        {
                            m_log.WriteLine("Resetting global relabeling to original labels for source '" + m_strSource + "'...");
                            m_db.ResetLabels(m_nProjectID, nSrcId);
                        }
 
                        if (rgMaps.Count > 0)
                            m_log.WriteLine("Starting global relabeling for source '" + m_strSource + "'...");
 
                        if (bHasConditions)
                        {
                            foreach (LabelMapping map in m_param.labelmapping_param.mapping)
                            {
                                m_db.SetLabelMapping(nSrcId, map);
                            }
                        }
                        else
                        {
                            foreach (KeyValuePair<int, List<int>> kv in rgMaps)
                            {
                                m_db.UpdateLabelMapping(nSrcId, kv.Key, kv.Value);
                            }
                        }
 
                        m_db.UpdateLabelCounts(m_nProjectID, nSrcId);
 
                        if (rgMaps.Count > 0)
                            m_log.WriteLine("Global relabeling completed for source '" + m_strSource + "'...");
 
                        if (m_param.labelmapping_param.label_boosts != null && m_param.labelmapping_param.label_boosts.Length > 0)
                        {
                            IXImageDatabase1 db = m_db as IXImageDatabase1;
                            bool bReloadImageSet = false;
 
                            if (db != null)
                            {
#warning ImageDatabase version 1 Only
                                string strLabelBoosts = db.GetLabelBoostsAsTextFromProject(m_nProjectID, nSrcId);
                                Dictionary<int, int> rgCounts = db.LoadLabelCounts(nSrcId);
                                string[] rgstrBoosts = m_param.labelmapping_param.label_boosts.Split(',');
                                Dictionary<int, int> rgBoostedLabelCounts = new Dictionary<int, int>();
                                double dfTotal = 0;
 
                                foreach (string strLabel in rgstrBoosts)
                                {
                                    int nLabel = int.Parse(strLabel);
 
                                    if (rgCounts.ContainsKey(nLabel))
                                    {
                                        int nCount = rgCounts[nLabel];
                                        rgBoostedLabelCounts.Add(nLabel, nCount);
                                        dfTotal += nCount;
                                    }
                                    else
                                    {
                                        rgBoostedLabelCounts.Add(nLabel, 0);
                                    }
                                }
 
                                db.DeleteLabelBoosts(m_nProjectID, nSrcId);
 
                                foreach (KeyValuePair<int, int> kv in rgCounts)
                                {
                                    double dfBoost = 0;
 
                                    if (rgBoostedLabelCounts.ContainsKey(kv.Key))
                                    {
                                        int nCount = rgBoostedLabelCounts.Count;
                                        dfBoost = (nCount == 0) ? 0 : (1.0 / (double)nCount);
                                    }
 
                                    db.AddLabelBoost(m_nProjectID, nSrcId, kv.Key, dfBoost);
                                }
 
                                string strNewLabelBoosts = db.GetLabelBoostsAsTextFromProject(m_nProjectID, nSrcId);
                                if (strNewLabelBoosts != strLabelBoosts)
                                    bReloadImageSet = true;
                            }
 
                            if (db != null)
                            {
                                string strNewLabelCounts = db.GetLabelCountsAsTextFromSourceId(nSrcId);
 
                                if (strNewLabelCounts != strLabelCounts || bReloadImageSet)
                                    db.ReloadImageSet(nSrcId);
                            }
 
                            m_log.WriteLine("WARNING: Label boosts are depreciated and soon to be removed.");
                        }
                        else
                        {
                            m_log.WriteLine("WARNING: ImageDatabase Version 2 currently does not support label mapping.");
                        }
                    }
                }
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nCount = colBottom[0].count();
 
            m_log.CHECK_EQ(nCount, colBottom[0].num, "The count should equal the number of items for the label blobs.");
            m_log.CHECK_EQ(nCount, colTop[0].count(), "The top and bottom should have the same number of items.");
 
            double[] rgBottom = convertD(colBottom[0].update_cpu_data());
 
            for (int i = 0; i < rgBottom.Length; i++)
            {
                int nLabel = m_param.labelmapping_param.MapLabel((int)rgBottom[i]);
                rgBottom[i] = nLabel;
 
                lock (m_syncActualMappedLabels)
                {
                    if (!m_rgActualMappedLabelCounts.ContainsKey(nLabel))
                        m_rgActualMappedLabelCounts.Add(nLabel, 1);
                    else
                        m_rgActualMappedLabelCounts[nLabel]++;
                }
            }
 
            colTop[0].mutable_cpu_data = convert(rgBottom);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
        }
    }
}